import librosa
import numpy as np
import keras
from keras.models import Sequential
from keras.layers import Dense
import math


print('读取训练信号\n')
for x in range(900):
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s1\S1_'+str(x+1)+'.wav', sr=None)
    y = y.reshape(1, len(y))
    if(x == 0):
        maleSignal = y
    else:
        maleSignal = np.c_[maleSignal, y]
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s16\S16_' + str(x + 1) + '.wav',sr=None)
    y = y.reshape(1, len(y))
    if (x == 0):
        femaleSignal = y
    else:
        femaleSignal = np.c_[femaleSignal, y]
'#33592400 38061900'
print(len(maleSignal[0]), len(femaleSignal[0]))
'#256点stft为1帧，33280000/256=130000帧语音'
maleSignal = np.array(maleSignal[0][0:33280000]).reshape(1, -1)
femaleSignal = np.array(femaleSignal[0][0:33280000]).reshape(1, -1)
print(len(maleSignal[0]))
MixedSignal = np.array(maleSignal[0]+femaleSignal[0]).reshape(1, -1)

print("读取测试信号\n")
for x in range(100):
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s1\S1_'+str(x+901)+'.wav', sr=None)
    y = y.reshape(1, len(y))
    if(x == 0):
        CMaleSignal = y
    else:
        CMaleSignal = np.c_[CMaleSignal, y]
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s16\S16_' + str(x + 901) + '.wav',sr=None)
    y = y.reshape(1, len(y))
    if (x == 0):
        CFemaleSignal = y
    else:
        CFemaleSignal = np.c_[CFemaleSignal, y]
'#3812600 4250350'
print(len(CMaleSignal[0]), len(CFemaleSignal[0]))
'#256点stft为1帧，3788800/256=14800帧语音'
CMaleSignal = np.array(CMaleSignal[0][0:3788800]).reshape(1, -1)
CFemaleSignal = np.array(CFemaleSignal[0][0:3788800]).reshape(1, -1)
print(len(CMaleSignal[0]), len(CMaleSignal[0]) == len(CFemaleSignal[0]))
CMixedSignal = np.array(CMaleSignal[0]+CFemaleSignal[0]).reshape(1, -1)


print("信号变换：")
# 男声信号
BmaleSignal = librosa.stft(maleSignal[0], n_fft=512)
BmaleSignalAngle = np.angle(BmaleSignal)
BmaleSignalAmplitude = np.abs(BmaleSignal)
TrainMaleSignal = BmaleSignalAmplitude.transpose()
# 女生信号
BFemaleSignal = librosa.stft(femaleSignal[0], n_fft=512)
BFemaleSignalAngle = np.angle(BFemaleSignal)
BFemaleSignalAmplitude = np.abs(BFemaleSignal)
TrainFemaleSignal = BFemaleSignalAmplitude.transpose()
# BmaleSignalLPS = librosa.power_to_db(BmaleSignalAmplitude**2, ref=1.0)
#混合信号
BMixedSignal = librosa.stft(MixedSignal[0], n_fft=512)
BMixedSignalAngle = np.angle(BmaleSignal)
BMixedSignalAmplitude = np.abs(BmaleSignal)
# BMixedSignalLPS = librosa.power_to_db(BmaleSignalAmplitude**2, ref=1.0)
TrainMixedSignal = BMixedSignalAmplitude.transpose()
#测试混合信号
TMixedSignal = librosa.stft(CMixedSignal[0], n_fft=512)
TMixedSignalAmplitude = np.abs(TMixedSignal)
TMixedSignalAngle = np.angle(TMixedSignal)
TestMixedSignal = TMixedSignalAmplitude.transpose()
print("信号变换完成")

print("开始训练：")
X_train = TrainMixedSignal
X_train = np.array(X_train)
Y_train = TrainMaleSignal
Y_train = np.array(Y_train)
Y_Ftrain = TrainFemaleSignal
Y_Ftrain = np.array(Y_Ftrain)
Y_Ctrain = np.c_[Y_train, Y_Ftrain]
X_test = TestMixedSignal
X_test = np.array(X_test)
model = Sequential()
model.add(Dense(1024, activation='relu', input_dim=257))
model.add(Dense(1024, activation='relu'))
model.add(Dense(2048, activation='sigmoid'))
model.add(Dense(257, activation='selu'))
sgd = keras.optimizers.SGD(lr=0.001, momentum=0.0, decay=0.0, nesterov=False)
# ADAM = keras.optimizers.Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=None, decay=0.0, amsgrad=False)
#NADAM =keras.optimizers.Nadam(lr=0.002, beta_1=0.9, beta_2=0.999)
#ADAMAX=keras.optimizers.Adamax(lr=0.002, beta_1=0.9, beta_2=0.999)
# adaBound = AdaBound(model.parameters(), lr=1e-3, final_lr=0.1)
#Adadelta=keras.optimizers.Adadelta(lr=1.0, rho=0.95)
model.compile(optimizer='adam', loss='mse', metrics=['mae'])
model.fit(X_train, Y_Ctrain, epochs=200, batch_size=256, verbose=2)
# model.evaluate(X_test, batch_size=256, verbose=0, steps=None)
Y_pred = model.predict(X_test, batch_size=256, verbose=0, steps=None)
print('训练完成\n')


print("语音输出开始：")
rMaleSignal, rFemaleSignal = np.split(Y_pred, [257], 1)
rMaleSignal = rMaleSignal.transpose()
rMaleSignal = rMaleSignal*np.power(math.e, TMixedSignalAngle*1j)
rFemaleSignal = rFemaleSignal.transpose()
rFemaleSignal = rFemaleSignal*np.power(math.e, TMixedSignalAngle*1j)
output1 = librosa.istft(rMaleSignal)
output2 = librosa.istft(rFemaleSignal)
librosa.output.write_wav(r".\speech\rebuild_male.wav", output1, sr=25000)
librosa.output.write_wav(r".\speech\rebuild_female.wav", output2, sr=25000)
librosa.output.write_wav(r".\speech\original_mixed.wav", CMixedSignal[0], sr=25000)
print("语音输出结束；")

